#!/usr/bin/env python3
"""
环境设置脚本
自动搭建AI动画生成系统的开发环境
"""

import os
import sys
import subprocess
import platform
import shutil
from pathlib import Path
import logging
import json
import argparse

logger = logging.getLogger(__name__)


class EnvironmentSetup:
    """环境设置器"""
    
    def __init__(self):
        self.project_root = Path(__file__).parent.parent
        self.system = platform.system().lower()
        self.python_version = sys.version_info
        
        # 检查Python版本
        if self.python_version < (3, 8):
            raise RuntimeError("需要Python 3.8或更高版本")
        
        # 设置日志
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s'
        )
    
    def check_conda(self) -> bool:
        """检查conda是否可用"""
        try:
            result = subprocess.run(['conda', '--version'], 
                                  capture_output=True, text=True, check=True)
            logger.info(f"找到conda: {result.stdout.strip()}")
            return True
        except (subprocess.CalledProcessError, FileNotFoundError):
            logger.warning("未找到conda，请先安装Anaconda或Miniconda")
            return False
    
    def check_gpu(self) -> bool:
        """检查GPU可用性"""
        try:
            import torch
            if torch.cuda.is_available():
                gpu_count = torch.cuda.device_count()
                gpu_name = torch.cuda.get_device_name(0)
                logger.info(f"找到GPU: {gpu_name} (数量: {gpu_count})")
                return True
            else:
                logger.warning("未找到CUDA GPU，将使用CPU模式")
                return False
        except ImportError:
            logger.warning("PyTorch未安装，无法检查GPU")
            return False
    
    def create_conda_environment(self, env_name: str = "ai-animation") -> bool:
        """创建conda环境"""
        try:
            logger.info(f"创建conda环境: {env_name}")
            
            # 检查环境是否已存在
            result = subprocess.run(['conda', 'env', 'list'], 
                                  capture_output=True, text=True, check=True)
            
            if env_name in result.stdout:
                logger.info(f"环境 {env_name} 已存在")
                return True
            
            # 创建环境
            subprocess.run(['conda', 'create', '-n', env_name, 'python=3.9', '-y'], 
                          check=True)
            
            logger.info(f"conda环境 {env_name} 创建成功")
            return True
            
        except subprocess.CalledProcessError as e:
            logger.error(f"创建conda环境失败: {e}")
            return False
    
    def install_dependencies(self, env_name: str = "ai-animation") -> bool:
        """安装依赖包"""
        try:
            logger.info("安装依赖包...")
            
            # 激活环境并安装依赖
            if self.system == "windows":
                activate_cmd = f"conda activate {env_name} && "
            else:
                activate_cmd = f"source activate {env_name} && "
            
            # 安装基础依赖
            packages = [
                "pytorch", "torchvision", "torchaudio", "cudatoolkit=11.8",
                "pip", "numpy", "scipy", "matplotlib", "pillow", "opencv-python",
                "tqdm", "requests", "huggingface-hub", "transformers", "diffusers"
            ]
            
            for package in packages:
                logger.info(f"安装 {package}")
                cmd = f"{activate_cmd}conda install -c pytorch -c conda-forge {package} -y"
                subprocess.run(cmd, shell=True, check=True)
            
            # 安装pip包
            pip_packages = [
                "accelerate", "xformers", "moviepy", "imageio", "imageio-ffmpeg",
                "scikit-image", "gradio", "streamlit", "fastapi", "uvicorn",
                "pydantic", "python-multipart", "aiofiles", "websockets",
                "redis", "celery", "flower", "prometheus-client", "psutil",
                "gputil", "tensorboard", "wandb", "mlflow", "optuna",
                "ray[tune]", "hydra-core", "omegaconf", "pytorch-lightning",
                "timm", "albumentations", "kornia", "einops", "safetensors",
                "sentencepiece", "protobuf", "onnx", "onnxruntime"
            ]
            
            for package in pip_packages:
                logger.info(f"pip安装 {package}")
                cmd = f"{activate_cmd}pip install {package}"
                subprocess.run(cmd, shell=True, check=True)
            
            logger.info("依赖包安装完成")
            return True
            
        except subprocess.CalledProcessError as e:
            logger.error(f"安装依赖包失败: {e}")
            return False
    
    def create_directories(self):
        """创建项目目录结构"""
        try:
            logger.info("创建项目目录结构...")
            
            directories = [
                "data/input",
                "data/output", 
                "data/temp",
                "models",
                "logs",
                "configs",
                "src/models",
                "src/utils",
                "src/api",
                "src/web",
                "scripts",
                "tests",
                "docs"
            ]
            
            for directory in directories:
                dir_path = self.project_root / directory
                dir_path.mkdir(parents=True, exist_ok=True)
                logger.info(f"创建目录: {directory}")
            
            logger.info("目录结构创建完成")
            
        except Exception as e:
            logger.error(f"创建目录结构失败: {e}")
            raise
    
    def create_config_files(self):
        """创建配置文件"""
        try:
            logger.info("创建配置文件...")
            
            # 模型配置文件
            model_config = {
                "text2video": {
                    "model": "runwayml/stable-diffusion-v1-5",
                    "device": "cuda",
                    "batch_size": 1,
                    "num_frames": 24,
                    "fps": 8
                },
                "style_transfer": {
                    "model": "stylegan3",
                    "device": "cuda",
                    "style_strength": 0.8
                },
                "video_editing": {
                    "model": "video-llama",
                    "device": "cuda",
                    "max_length": 512
                }
            }
            
            config_file = self.project_root / "configs" / "models.yaml"
            config_file.parent.mkdir(parents=True, exist_ok=True)
            
            import yaml
            with open(config_file, 'w', encoding='utf-8') as f:
                yaml.dump(model_config, f, default_flow_style=False, allow_unicode=True)
            
            # 系统配置文件
            system_config = {
                "gpu": {
                    "memory_limit": "8GB",
                    "precision": "fp16"
                },
                "storage": {
                    "input_dir": "data/input",
                    "output_dir": "data/output",
                    "temp_dir": "data/temp",
                    "max_file_size": "2GB"
                },
                "api": {
                    "host": "0.0.0.0",
                    "port": 8000,
                    "workers": 4,
                    "timeout": 300
                }
            }
            
            system_config_file = self.project_root / "configs" / "system.yaml"
            with open(system_config_file, 'w', encoding='utf-8') as f:
                yaml.dump(system_config, f, default_flow_style=False, allow_unicode=True)
            
            logger.info("配置文件创建完成")
            
        except Exception as e:
            logger.error(f"创建配置文件失败: {e}")
            raise
    
    def download_models(self):
        """下载预训练模型"""
        try:
            logger.info("开始下载预训练模型...")
            
            # 运行模型下载脚本
            download_script = self.project_root / "scripts" / "download_models.py"
            
            if download_script.exists():
                subprocess.run([sys.executable, str(download_script), "--status"], check=True)
            else:
                logger.warning("模型下载脚本不存在")
            
        except subprocess.CalledProcessError as e:
            logger.error(f"下载模型失败: {e}")
        except Exception as e:
            logger.error(f"下载模型时出错: {e}")
    
    def create_startup_scripts(self, env_name: str = "ai-animation"):
        """创建启动脚本"""
        try:
            logger.info("创建启动脚本...")
            
            # Windows批处理脚本
            if self.system == "windows":
                bat_script = self.project_root / "start.bat"
                bat_content = f"""@echo off
echo 启动AI动画生成系统...
call conda activate {env_name}
python main.py web_ui
pause
"""
                with open(bat_script, 'w', encoding='utf-8') as f:
                    f.write(bat_content)
                
                logger.info(f"创建启动脚本: {bat_script}")
            
            # Unix shell脚本
            else:
                sh_script = self.project_root / "start.sh"
                sh_content = f"""#!/bin/bash
echo "启动AI动画生成系统..."
source activate {env_name}
python main.py web_ui
"""
                with open(sh_script, 'w', encoding='utf-8') as f:
                    f.write(sh_content)
                
                # 设置执行权限
                os.chmod(sh_script, 0o755)
                logger.info(f"创建启动脚本: {sh_script}")
            
            logger.info("启动脚本创建完成")
            
        except Exception as e:
            logger.error(f"创建启动脚本失败: {e}")
            raise
    
    def run_tests(self):
        """运行测试"""
        try:
            logger.info("运行系统测试...")
            
            # 检查Python包导入
            test_imports = [
                "torch", "torchvision", "numpy", "cv2", "PIL", 
                "transformers", "diffusers", "gradio"
            ]
            
            for package in test_imports:
                try:
                    __import__(package)
                    logger.info(f"✓ {package} 导入成功")
                except ImportError as e:
                    logger.warning(f"✗ {package} 导入失败: {e}")
            
            # 检查GPU
            if self.check_gpu():
                logger.info("✓ GPU可用")
            else:
                logger.warning("✗ GPU不可用，将使用CPU模式")
            
            logger.info("系统测试完成")
            
        except Exception as e:
            logger.error(f"运行测试失败: {e}")
    
    def setup_complete(self, env_name: str = "ai-animation"):
        """完成设置"""
        try:
            logger.info("=" * 50)
            logger.info("AI动画生成系统环境设置完成！")
            logger.info("=" * 50)
            
            print("\n🎉 环境设置完成！")
            print("\n📋 下一步操作:")
            print("1. 激活conda环境:")
            if self.system == "windows":
                print(f"   conda activate {env_name}")
            else:
                print(f"   source activate {env_name}")
            
            print("\n2. 启动Web界面:")
            print("   python main.py web_ui")
            
            print("\n3. 或使用启动脚本:")
            if self.system == "windows":
                print("   start.bat")
            else:
                print("   ./start.sh")
            
            print("\n4. 查看帮助:")
            print("   python main.py --help")
            
            print("\n📁 项目结构:")
            print("   data/          - 数据目录")
            print("   models/        - 模型目录")
            print("   src/           - 源代码")
            print("   configs/       - 配置文件")
            print("   scripts/       - 脚本文件")
            
            print("\n🔧 常用命令:")
            print("   python main.py text2video --prompt '描述'")
            print("   python main.py style_transfer --input video.mp4 --style cartoon")
            print("   python main.py video_edit --input video.mp4 --operations trim resize")
            
            print("\n📚 更多信息请查看 README.md")
            print("=" * 50)
            
        except Exception as e:
            logger.error(f"完成设置时出错: {e}")
    
    def setup(self, env_name: str = "ai-animation", skip_models: bool = False):
        """执行完整设置"""
        try:
            logger.info("开始AI动画生成系统环境设置...")
            
            # 1. 检查系统要求
            logger.info("检查系统要求...")
            if not self.check_conda():
                raise RuntimeError("需要安装conda")
            
            self.check_gpu()
            
            # 2. 创建conda环境
            if not self.create_conda_environment(env_name):
                raise RuntimeError("创建conda环境失败")
            
            # 3. 安装依赖
            if not self.install_dependencies(env_name):
                raise RuntimeError("安装依赖失败")
            
            # 4. 创建目录结构
            self.create_directories()
            
            # 5. 创建配置文件
            self.create_config_files()
            
            # 6. 下载模型（可选）
            if not skip_models:
                self.download_models()
            
            # 7. 创建启动脚本
            self.create_startup_scripts(env_name)
            
            # 8. 运行测试
            self.run_tests()
            
            # 9. 完成设置
            self.setup_complete(env_name)
            
            return True
            
        except Exception as e:
            logger.error(f"环境设置失败: {e}")
            return False


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="AI动画生成系统环境设置")
    parser.add_argument("--env-name", type=str, default="ai-animation", 
                       help="conda环境名称")
    parser.add_argument("--skip-models", action="store_true", 
                       help="跳过模型下载")
    parser.add_argument("--check-only", action="store_true", 
                       help="仅检查环境")
    
    args = parser.parse_args()
    
    try:
        setup = EnvironmentSetup()
        
        if args.check_only:
            # 仅检查环境
            print("检查系统环境...")
            setup.check_conda()
            setup.check_gpu()
            setup.run_tests()
        else:
            # 完整设置
            success = setup.setup(args.env_name, args.skip_models)
            if not success:
                sys.exit(1)
        
    except Exception as e:
        logger.error(f"设置失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main() 